from langchain_core.messages import AIMessage
from langchain_ollama import ChatOllama
from langchain_core.prompts import ChatPromptTemplate

from langchain_core.output_parsers import StrOutputParser
messages = [
    (
        "system",
        "你是一个优秀的StableDiffusion提示词专家，你会编写非常生动的提示词。非常擅长FlUX提示词的编写，用短句的方式表达出来。注意同时给出中英双语的答案。"
    ),
    (
        "human",
        "请你描述一下宽阔的大海上的一艘巨轮正在航行，阳光明媚，海鸟飞舞，在海平面的远处有一个若隐若现的巨大海怪。"
    )
]

# 应用上面的 messages
prompt = ChatPromptTemplate.from_messages(messages)


llm = ChatOllama(
    model="llama3:8b-instruct-fp16",
    temperature=0.1,
)

parser = StrOutputParser()
chain = prompt | llm | parser


ai_msg = response = chain.invoke({})

print(ai_msg)
